from typing import Union, Dict

import numpy as np


class ChannelShuffle:
    """
    Randomly shuffle channels in EEG data matrix.

    Part of the CMSAugment policy proposed in [1]_

    Parameters
    ----------
    prob: float
        Float between 0 and 1 setting the probability of including the channel
        in the set of permuted channels.
    seed: int | np.random.Generator, optional
        Seed to be used to instantiate numpy random number generator instance.
        Used to sample which channels to shuffle and to carry the shuffle.
        Defaults to None.

    References
    ----------
    .. [1] Saeed, A., Grangier, D., Pietquin, O., & Zeghidour, N. (2020).
       Learning from Heterogeneous EEG Signals with Differentiable Channel
       Reordering. arXiv preprint arXiv:2010.13694.
    """
    def __init__(self, prob: float, seed=None, keys=None) -> None:
        self.prob = prob
        self.rng = np.random.default_rng(seed)
        self.keys = keys

    def _pick_channels_randomly(self, n_channels, p_no_shuffle):
        """Helper function to pick channels randomly."""
        mask = self.rng.random(n_channels) > p_no_shuffle
        return mask

    def _make_permutation_indices(self, n_channels, mask):
        """Helper function to create permutation indices for the channels."""
        perm_indices = np.arange(n_channels)
        shuffle_indices = self.rng.permutation(np.where(mask)[0])
        perm_indices[mask] = shuffle_indices
        return perm_indices
    
    def _apply(self, sample: np.ndarray) -> np.ndarray:
        """
        Shuffle the channels of the provided sample.

        :param sample: numpy array of shape (channels, ...)
        :return: numpy array with shuffled channels
        """
        n_channels = sample.shape[0]
        if self.rng.random() < self.prob:
            permutation = self.rng.permutation(n_channels)
            return sample[permutation]
        return sample

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                if key in data:
                    data[key] = self._apply(data[key])
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data